Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

Differential Revision: D83924526

Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

Differential Revision: D83924526

Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

Differential Revision: D83924526

Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

Reviewed By: spmex

Differential Revision: D83924526

Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184} 
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242} 
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281} 
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334} 
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765} 

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526

Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184} 
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242} 
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281} 
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334} 
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765} 

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526

Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184} 
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242} 
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281} 
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334} 
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765} 

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526

Summary:

# context
* add benchmark for `all_to_all_single` with `async_op` option. 
* the comms uses a different cuda stream (comms stream) so it's non-blocking for the followed operations (on main cuda stream)
* of course the comms results (pre-allocated output) are not valid until the comms' done (the pre-check fails)
* when there's data dependency on the comms' output, user's need to call `req.wait()` explicitly, so that the main cuda stream wait on the comms stream
NOTE: the `req.wait()` call is non-blocking on the CPU side.

# ref
* [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)
* [sync and async comms](https://docs.pytorch.org/docs/stable/distributed.html#collective-functions)
* [CUDA semantics](https://docs.pytorch.org/docs/stable/notes/cuda.html)

# optimization
* the data validation is done on the device side, and very often the validation result is needed from the cpu side
* in a previously example (a2a-sync), the assertion needs the `checks` on the cpu side, so it's blocking the cpu execution (see below)
 {F1982527184} 
* actually the `checks` is available at the gpu side once the "comms validation" is done, and a device-to-host data transfer can be initiated. however, this device-to-host data transfer is blocking the following cpu execution (i.e., "irrelevant compute")
 {F1982527242} 
* we tried to use `non_blocking=True` in the copy_to, and the results are very interesting:
all the cpu executions are done very early, this is because the cpu has no idea about the completeness of non-blocking device-to-host data transfer, it will just go ahead executing, which of coruse causes assertion error (changed to `print(checks)` to bypass the assertion in order to generate the trace)
 {F1982527281} 
* the proper way of doing this is to use cuda.event as in this diff, so that the assert only needs to wait until the data (`checks`) becomes available on the cpu side.
 {F1982527334} 
* as for the async all_to_all_single case, the 2nd batch (CPU) starts right after the 1st assertion, which is done right after the device-to-host data transfer.
{F1982527765} 

# results
* [a2a sync trace](https://drive.google.com/file/d/1xI-qlI7V6zbmcbuQGyZgqFC3ZMY1scro/view?usp=sharing)
* [a2a async trace](https://drive.google.com/file/d/1eboZ01quSZjKBtBcdu4vXX9zyokzohqi/view?usp=sharing)

Reviewed By: spmex

Differential Revision: D83924526

Summary:
# TL;DR
* A new `DeviceToHostTensorAwaitable` class is available to wrap the device-to-host data transfer, and defers the `cudaEventSync` call until the data is really used on the host.
* It aims at helping sync-point removal in training optimization which often suffers from cpu-blocking sync points.

# why awaitable
* as shown in the following diagram, a comms op is often better to overlap with another (irrelevant) compute op to better utilize the device capability
* the idea is to **defer** the `wait()` call until running the function that uses the result from the comm op
* a convenient way to achieve this "deferring" behavior is to use the `lazy_awaitable` concept, which is already [implemented in torchrec](https://github.com/meta-pytorch/torchrec/blob/main/torchrec/distributed/types.py#L368)
* diagram of (lazy_)awaitable in torchrec
 {F1982900178}

# why device-to-host transfer
* there are scenarios that the on-device data is needed from the host side, such as metrics logging and data-dependent shape operation.
* those pattern creates a device-to-host sync (data transfer) that often blocks the cpu execution, and the correct implementation (with `.to(non_blocking=True)` and cuda event: [PR 3436](meta-pytorch#3436)) usually spans across multiple code domain making it difficult to optimize. 
* here we borrow the `LazyAwaitable` concept for the device-side comms and wrap the (1) non-blocking device-to-host data transfer, and (2) `cuda_event.wait()` inside a `DeviceToHostTensorAwaitable` class for better user experience.
* diagram of lazy_awaitable for device-to-host data transfer
 {F1982900233} 

# results
* the "comms check" result is on device and is needed for validation (host-side assertion)
* the `DeviceToHostTensorAwaitable.wait()` **defer** the cudaEventSync until the very end where the result is really needed by host. 
* You can see the post-comms computes are scheduled before the assertion on the host side.
{F1982900468} 

NOTE: in this version of implementation we don't use a separate stream (as shown in the diagram above) for the non-blocking device-to-host data transfer because usually the data volume is relatively small.

Differential Revision: D85211205

Summary:
Pull Request resolved: #3477
workplace post: https://fb.workplace.com/groups/811751593969209/permalink/1285164823294548/

# TL;DR
* A new `DeviceToHostTensorAwaitable` class is available to wrap the device-to-host data transfer, and defers the `cudaEventSync` call until the data is really used on the host.
* It aims at helping sync-point removal in training optimization which often suffers from cpu-blocking sync points.

# why awaitable
* as shown in the following diagram, a comms op is often better to overlap with another (irrelevant) compute op to better utilize the device capability
* the idea is to **defer** the `wait()` call until running the function that uses the result from the comm op
* a convenient way to achieve this "deferring" behavior is to use the `lazy_awaitable` concept, which is already [implemented in torchrec](https://github.com/meta-pytorch/torchrec/blob/main/torchrec/distributed/types.py#L368)
* diagram of (lazy_)awaitable in torchrec
 {F1982900178}

# why device-to-host transfer
* there are scenarios that the on-device data is needed from the host side, such as metrics logging and data-dependent shape operation.
* those pattern creates a device-to-host sync (data transfer) that often blocks the cpu execution, and the correct implementation (with `.to(non_blocking=True)` and cuda event: [PR 3436](#3436)) usually spans across multiple code domain making it difficult to optimize.
* here we borrow the `LazyAwaitable` concept for the device-side comms and wrap the (1) non-blocking device-to-host data transfer, and (2) `cuda_event.wait()` inside a `DeviceToHostTensorAwaitable` class for better user experience.
* diagram of lazy_awaitable for device-to-host data transfer
 {F1982900233}

# results
* the "comms check" result is on device and is needed for validation (host-side assertion)
* the `DeviceToHostTensorAwaitable.wait()` **defer** the cudaEventSync until the very end where the result is really needed by host.
* You can see the post-comms computes are scheduled before the assertion on the host side.
{F1982900468}

NOTE: in this version of implementation we don't use a separate stream (as shown in the diagram above) for the non-blocking device-to-host data transfer because usually the data volume is relatively small.
{F1982901286}

Reviewed By: spmex

Differential Revision: D85211205

fbshipit-source-id: 41d03230dd9b190085545cfb76192d59375646c4